package digitRecognitionProblem.reduceConnections;

import genetic_algorithm.Chromosome;
import genetic_algorithm.Crossover;
import genetic_algorithm.FitnessFunction;
import genetic_algorithm.GeneticAlgorithm;
import genetic_algorithm.Mutation;
import genetic_algorithm.Population;
import genetic_algorithm.Selection;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.PrintStream;
import java.util.LinkedList;
import java.util.List;

import mlp.MlpExamples;
import mlp.SparseMlp;
import utils.GAretVal;
import utils.RandomGenerator;

/**
 * Runs a genetic algorithm that chooses which edges, of a previously-trained
 * neural network of a certain digit, should be trimmed.
 * At the end of the run, a BitSet representing the edges is serialized.
 * ATTENTION: before running, the activation function of mlp.Neuron.java must be set to tanh
 */
public class ReduceConMain {

	// indices for main arguments
	private enum ArgNums {
		MLP_DIRECTORY, TEST_FILE, START_DIGIT, END_DIGIT, NUM_EPOCHS
	}
	
	private static final int MAX_GENS = 100; // maximal number of generations to run
	private static final int CROSSOVER_MODE = ReduceConCrossover.SINGLE_POINT; // the mode of DigitRecognitionCrossover;
	private static final int SELECTION_MODE = ReduceConSelection.ROULETTE; // selection phase mode
	private static final int TOURNAMENT_SIZE = 5; // how many chromosomes to include in tournament (if such mode is chosen for selection)
	private static final int POP_SIZE = 100; // number of chromosomes in the population
	private static final double CROSSOVER_RATE = 0.9; // chance to perform crossover
	private static final float LEARNING_RATE = 0.01f; // learning rate for back propagation
	private static final float MOMENTUM = 0.2f; // momentum weight for back propagation
	private static final double MUTATION_RATE = 0.25; // chance to perform mutation
	private static int NUM_EPOCHS = 40; // number of iterations to run back-propagation after creating a new chromosome
	public static final String TRAIN_FILE = "train_reduce.txt"; // path to train file	
	private static final int TRAIN_AMOUNT = 200; // how many train examples to use
	private static String TEST_FILE = "validate1_reduce.txt"; // path to test file
	private static boolean REMOVE_PARENTS = false; // if chromosomes can be selected only once
	
	public static int DIGIT = 0; // digit to learn
		
	public static void main(String[] args) {

		// parse start and end digits
		int startDigit = Integer.parseInt(args[ArgNums.START_DIGIT.ordinal()]);
		int endDigit = Integer.parseInt(args[ArgNums.END_DIGIT.ordinal()]);
		
		// redirect output to file
		try {
			File file = new File("Details_" + startDigit + "_to_" + endDigit + ".txt");
			FileOutputStream fos;
			fos = new FileOutputStream(file);
			PrintStream ps = new PrintStream(fos);
			System.setOut(ps);
			System.setErr(ps);
		} catch (FileNotFoundException e1) {
			e1.printStackTrace();
		}
		
		// parse test file
		TEST_FILE = args[ArgNums.TEST_FILE.ordinal()];		
		
		RandomGenerator.showSeed();
		
		// set global parameters for neural network
		SparseMlp.LEARNING_RATE = LEARNING_RATE;
		SparseMlp.MOMENTUM = MOMENTUM;
		SparseMlp.NUMBER_EXAMPLE = TRAIN_AMOUNT;
		
		NUM_EPOCHS = Integer.parseInt(args[ArgNums.NUM_EPOCHS.ordinal()]);

		for (int digit = startDigit ; digit < endDigit + 1 ; ++digit) {
			
			System.out.println("current digit: " + DIGIT);
		
			// set current digit
			DIGIT = digit;
			
			// set path to serialized neural network
			ReduceConFitness.MLP_FILE = args[ArgNums.MLP_DIRECTORY.ordinal()]
					+ "/" + DIGIT + "/mlp_0_rec_" + DIGIT + ".txt";

			// extract the trained network
			SparseMlp mlp = null;
			try {
				FileInputStream readFile = new FileInputStream(new File(ReduceConFitness.MLP_FILE));
				ObjectInputStream fileStream = new ObjectInputStream(readFile);
				mlp = (SparseMlp)fileStream.readObject();
				fileStream.close();
			} catch (Exception e) {
				e.printStackTrace();
			}
			
			// set number of neurons in each layer
			SparseMlp.HIDDEN_LAYYER = mlp.getLayer(0).size();
			SparseMlp.OUTPUT_LAYYER = mlp.getLayer(1).size();
			
			// parse test examples
			MlpExamples test = new MlpExamples(TEST_FILE, DIGIT);
			
			System.out.println("train and test examples are loaded");
			
			// allocate objects to implement each phase of the genetic algorithm
			FitnessFunction fitnessFunction = new ReduceConFitness(TRAIN_AMOUNT, test, NUM_EPOCHS);
			Selection selector = new ReduceConSelection(REMOVE_PARENTS, SELECTION_MODE, TOURNAMENT_SIZE);
			Crossover crossover = new ReduceConCrossover(CROSSOVER_MODE);
			Mutation mutation = new ReduceConMutation();
			
			System.out.println("initializing population");
			
			// create initial population
			List<Chromosome> chromList = new LinkedList<Chromosome>();
			for (int i = 0; i < POP_SIZE; ++i) {
				chromList.add(new ReduceConChromosome());
			}
			Population population = new Population(chromList, fitnessFunction);
			
			System.out.println("population is set");
			
			// initialize the genetic algorithm
			GeneticAlgorithm GA = new GeneticAlgorithm(fitnessFunction, crossover,
					selector, mutation, CROSSOVER_RATE, MUTATION_RATE);
			
			// find optimal solution
			GAretVal result = GA.execute(population, MAX_GENS, true, 6);
			
			// generate graph of the execution data
			result.createsCsvFiles("reduce_con_digit_"+DIGIT);
			
			// serialize the bit set
			FileOutputStream writeFile;
			try {
				writeFile = new FileOutputStream("edge_switch_rec_" + DIGIT + ".txt");
				ObjectOutputStream fileStream = new ObjectOutputStream(writeFile);
				fileStream.writeObject(((ReduceConChromosome) result.getOptSolution()).getData());
				fileStream.close();
			} catch (FileNotFoundException e) {
				e.printStackTrace();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		
	}
}
